import tiktoken

def count_deepseek_tokens_approx(text, model_name="gpt-4"):
    """
    近似计算 DeepSeek Token 数量的函数
    
    参数:
    text (str): 需要计算 token 的文本
    model_name (str): 用于近似计算的模型名称，默认为 "gpt-4"
    
    返回:
    int: 文本的近似 token 数量
    list: token ID 列表（可选，用于调试）
    """
    try:
        # 获取编码器
        encoding = tiktoken.encoding_for_model(model_name)
        
        # 编码文本为 token IDs
        token_ids = encoding.encode(text)
        
        # 返回 token 数量和 token IDs
        return len(token_ids), token_ids
        
    except KeyError:
        # 如果指定的模型不存在，使用 cl100k_base 编码（GPT-4 使用的编码）
        encoding = tiktoken.get_encoding("cl100k_base")
        token_ids = encoding.encode(text)
        return len(token_ids), token_ids

def tokenize_and_display(text, model_name="gpt-4"):
    """
    将文本分词并显示详细信息
    
    参数:
    text (str): 需要分词的文本
    model_name (str): 用于近似计算的模型名称
    """
    # 计算 token 数量
    token_count, token_ids = count_deepseek_tokens_approx(text, model_name)
    
    # 获取编码器以进行分词
    try:
        encoding = tiktoken.encoding_for_model(model_name)
    except KeyError:
        encoding = tiktoken.get_encoding("cl100k_base")
    
    # 将 token IDs 解码为 token 字符串
    tokens = [encoding.decode_single_token_bytes(token_id).decode('utf-8', errors='replace') 
              for token_id in token_ids]
    
    # 打印结果
    print(f"文本: {text}")
    print(f"近似 Token 数量: {token_count}")
    print("Token 拆分详情:")
    for i, (token_id, token_str) in enumerate(zip(token_ids, tokens)):
        print(f"  Token {i+1}: ID={token_id}, 内容='{token_str}'")
    
    return token_count, tokens

# 示例使用
if __name__ == "__main__":
    # 测试文本
    test_texts = [
        "Hello, world!",
        "你好，世界！",
        "这是一段混合中文和英文的文本。This is a mixed Chinese and English text.",
        "DeepSeek-V3 是一个强大的AI模型。"
    ]
    
    for text in test_texts:
        print("=" * 50)
        token_count, tokens = tokenize_and_display(text)
        print()